import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LinearRegression, Ridge, Lasso
# 正确的导入方式
from sklearn.tree import DecisionTreeRegressor
from sklearn.ensemble import RandomForestRegressor, GradientBoostingRegressor
from sklearn.svm import SVR
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
import warnings
warnings.filterwarnings('ignore')

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'Arial Unicode MS', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False

class RevenuePredictionAnalysis:
    def __init__(self):
        self.df = None
        self.X_train = None
        self.X_test = None
        self.y_train = None
        self.y_test = None
        self.models = {}
        self.scaler = StandardScaler()
        self.results = []
        self.feature_importance = {}
    
    def visualize_feature_importance(self, model_name, importances):
        """可视化特征重要性"""
        if not importances:
            return
        
        plt.figure(figsize=(12, 8))
        sorted_importances = sorted(importances.items(), key=lambda x: x[1], reverse=True)[:10]
        features = [item[0] for item in sorted_importances]
        values = [item[1] for item in sorted_importances]
        
        # 使用渐变色条形图
        colors = plt.cm.YlOrRd(np.linspace(0.2, 0.8, len(features)))
        bars = plt.barh(features, values, color=colors)
        
        # 添加数值标签
        for bar in bars:
            width = bar.get_width()
            plt.text(width + 0.01, bar.get_y() + bar.get_height()/2, f'{width:.4f}',
                     ha='left', va='center', fontweight='bold')
        
        plt.title(f'{model_name} - 特征重要性分析', fontsize=16, pad=20)
        plt.xlabel('重要性得分', fontsize=12, labelpad=10)
        plt.grid(axis='x', linestyle='--', alpha=0.7)
        plt.tight_layout()
        plt.savefig(f'{model_name}_特征重要性.png', dpi=300, bbox_inches='tight')
        plt.close()
        print(f'特征重要性图片已生成: {model_name}_特征重要性.png')
        
    def visualize_correlation(self):
        """可视化相关性热力图"""
        corr_df = self.df.select_dtypes(include=[np.number]).corr()
        plt.figure(figsize=(14, 12))
        
        # 使用seaborn热力图
        mask = np.triu(np.ones_like(corr_df, dtype=bool))
        cmap = sns.diverging_palette(230, 20, as_cmap=True)
        sns.heatmap(corr_df, mask=mask, cmap=cmap, annot=True, fmt='.2f', 
                   square=True, linewidths=.5, cbar_kws={"shrink": .8})
        
        plt.title('财务指标相关性热力图', fontsize=16, pad=20)
        plt.tight_layout()
        plt.savefig('相关性热力图.png', dpi=300, bbox_inches='tight')
        plt.close()
        print('相关性热力图已生成: 相关性热力图.png')
        
    def visualize_model_performance(self):
        """可视化模型性能对比"""
        if not self.results:
            return
        
        df = pd.DataFrame(self.results)
        df = df.sort_values('R²', ascending=False)
        
        # 创建子图
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 6))
        
        # RMSE对比
        colors = plt.cm.Blues(np.linspace(0.4, 0.8, len(df)))
        bars1 = ax1.bar(df['模型'], df['RMSE'], color=colors)
        ax1.set_title('模型RMSE对比', fontsize=14, pad=10)
        ax1.set_ylabel('RMSE', fontsize=12)
        ax1.set_xticklabels(df['模型'], rotation=45, ha='right')
        ax1.grid(axis='y', linestyle='--', alpha=0.7)
        
        # R²对比
        colors = plt.cm.Greens(np.linspace(0.4, 0.8, len(df)))
        bars2 = ax2.bar(df['模型'], df['R²'], color=colors)
        ax2.set_title('模型R²对比', fontsize=14, pad=10)
        ax2.set_ylabel('R²', fontsize=12)
        ax2.set_xticklabels(df['模型'], rotation=45, ha='right')
        ax2.set_ylim(0, 1)
        ax2.grid(axis='y', linestyle='--', alpha=0.7)
        
        # 添加数值标签
        for bar in bars1:
            height = bar.get_height()
            ax1.text(bar.get_x() + bar.get_width()/2., height + 0.01,
                    f'{height:.4f}', ha='center', va='bottom')
                    
        for bar in bars2:
            height = bar.get_height()
            ax2.text(bar.get_x() + bar.get_width()/2., height + 0.01,
                    f'{height:.4f}', ha='center', va='bottom')
        
        plt.tight_layout()
        plt.savefig('模型性能对比.png', dpi=300, bbox_inches='tight')
        plt.close()
        print('模型性能对比图片已生成: 模型性能对比.png')
    
    def load_data(self, file_path):
        """加载数据"""
        self.df = pd.read_excel(file_path)
        print(f"数据加载成功，共 {self.df.shape[0]} 条记录，{self.df.shape[1]} 个特征")
    
    def data_exploration(self):
        """数据探索"""
        print("\n=== 数据基本信息 ===")
        print(self.df.info())
        
        print("\n=== 描述性统计 ===")
        print(self.df.describe())
        
        # 检查缺失值
        print("\n=== 缺失值检查 ===")
        print(self.df.isnull().sum())
    
    def feature_engineering(self):
        """特征工程"""
        # 复制数据以避免修改原始数据
        df_processed = self.df.copy()
        
        # 创建新特征
        df_processed['毛利率'] = (df_processed['营业收入'] - df_processed['营业成本']) / df_processed['营业收入']
        df_processed['资产负债率'] = df_processed['负债合计'] / df_processed['资产总计']
        df_processed['净资产收益率'] = df_processed['归属于母公司所有者权益合计'] / df_processed['营业收入']
        df_processed['现金流量比率'] = df_processed['经营活动产生的现金流量净额'] / df_processed['营业收入']
        df_processed['投资活动比率'] = df_processed['投资活动产生的现金流量净额'] / df_processed['营业收入']
        df_processed['筹资活动比率'] = df_processed['筹资活动产生的现金流量净额'] / df_processed['营业收入']
        
        # 处理无穷大和NaN值
        df_processed = df_processed.replace([np.inf, -np.inf], np.nan)
        df_processed = df_processed.fillna(0)
        
        # 选择特征变量
        features = ['营业成本', '资产总计', '负债合计', '归属于母公司所有者权益合计', 
                   '所有者权益合计', '经营活动产生的现金流量净额', '投资活动产生的现金流量净额', 
                   '筹资活动产生的现金流量净额', '现金及现金等价物净增加额', '毛利率', 
                   '资产负债率', '净资产收益率', '现金流量比率', '投资活动比率', '筹资活动比率']
        
        X = df_processed[features]
        y = df_processed['营业收入']
        
        # 划分训练集和测试集
        self.X_train, self.X_test, self.y_train, self.y_test = train_test_split(
            X, y, test_size=0.2, random_state=42)
        
        # 标准化特征
        self.X_train_scaled = self.scaler.fit_transform(self.X_train)
        self.X_test_scaled = self.scaler.transform(self.X_test)
        
        print("\n=== 特征工程完成 ===")
        print(f"训练集大小: {self.X_train.shape[0]}, 测试集大小: {self.X_test.shape[0]}")
        print(f"特征数量: {self.X_train.shape[1]}")
    
    def train_models(self):
        """训练多种模型"""
        print("\n=== 模型训练开始 ===")
        
        # 定义要训练的模型
        model_list = [
            ('线性回归', LinearRegression()),
            ('Ridge回归', Ridge(alpha=1.0)),
            ('Lasso回归', Lasso(alpha=1.0)),
            ('决策树回归', DecisionTreeRegressor(max_depth=10, random_state=42)),
            ('随机森林回归', RandomForestRegressor(n_estimators=100, max_depth=10, random_state=42)),
            ('梯度提升回归', GradientBoostingRegressor(n_estimators=100, max_depth=5, random_state=42)),
            ('SVR', SVR(C=1.0, kernel='rbf'))
        ]
        
        for name, model in model_list:
            print(f"\n训练模型: {name}")
            if name == 'SVR':
                model.fit(self.X_train_scaled, self.y_train)
            else:
                model.fit(self.X_train, self.y_train)
            
            self.models[name] = model
            
            # 预测
            if name == 'SVR':
                y_pred = model.predict(self.X_test_scaled)
            else:
                y_pred = model.predict(self.X_test)
            
            # 评估模型
            mse = mean_squared_error(self.y_test, y_pred)
            rmse = np.sqrt(mse)
            mae = mean_absolute_error(self.y_test, y_pred)
            r2 = r2_score(self.y_test, y_pred)
            
            self.results.append({
                '模型': name,
                'MSE': mse,
                'RMSE': rmse,
                'MAE': mae,
                'R²': r2
            })
            
            print(f"  MSE: {mse:.4f}")
            print(f"  RMSE: {rmse:.4f}")
            print(f"  MAE: {mae:.4f}")
            print(f"  R²: {r2:.4f}")
            
            # 计算特征重要性（如果模型支持）
            if hasattr(model, 'feature_importances_'):
                self.feature_importance[name] = dict(zip(self.X_train.columns, model.feature_importances_))
            elif hasattr(model, 'coef_'):
                self.feature_importance[name] = dict(zip(self.X_train.columns, abs(model.coef_)))
    
    def analyze_results(self):
        """分析模型结果"""
        print("\n=== 模型性能对比 ===")
        results_df = pd.DataFrame(self.results)
        print(results_df.sort_values('R²', ascending=False))
        
        # 找出最佳模型
        best_model = results_df.loc[results_df['R²'].idxmax()]
        print(f"\n最佳模型: {best_model['模型']}")
        print(f"R²值: {best_model['R²']:.4f}")
        
        return best_model['模型']
    
    def analyze_feature_importance(self, best_model_name):
        """分析特征重要性"""
        print(f"\n=== {best_model_name} 的特征重要性 ===")
        
        if best_model_name in self.feature_importance:
            importances = self.feature_importance[best_model_name]
            sorted_importances = sorted(importances.items(), key=lambda x: x[1], reverse=True)
            
            for feature, importance in sorted_importances[:10]:
                print(f"{feature}: {importance:.4f}")
            
            return sorted_importances
        else:
            print("该模型不支持特征重要性分析")
            return []
    
    def generate_report(self):
        """不生成HTML报告（根据要求）"""
        pass

def main():
    # 创建分析对象
    analysis = RevenuePredictionAnalysis()
    
    # 加载数据
    analysis.load_data('上市公司营收数据.xlsx')
    
    # 数据探索
    analysis.data_exploration()
    
    # 特征工程
    analysis.feature_engineering()
    
    # 训练模型
    analysis.train_models()
    
    # 分析结果
    best_model = analysis.analyze_results()
    
    # 特征重要性分析
    analysis.analyze_feature_importance(best_model)
    
    # 可视化结果
    analysis.visualize_model_performance()
    analysis.visualize_correlation()
    
    # 不生成HTML报告（根据要求）
    # analysis.generate_report()

if __name__ == "__main__":
    main()